import frappe
from frappe import _
from frappe.utils import flt
from frappe.query_builder.functions import Count,Sum, Cast
from frappe.query_builder import Case, Criterion
from erpnext.manufacturing.doctype.work_order.work_order import WorkOrder
from erpnext.stock.doctype.stock_entry.stock_entry import StockEntry
from erpnext.stock.doctype.pick_list.pick_list import PickList


class CustomWorkOrder(WorkOrder):
    def update_transferred_qty_for_required_items(self):
        """update transferred qty from submitted stock entries for that item against
        the work order"""

        se = frappe.qb.DocType('Stock Entry')
        sed = frappe.qb.DocType('Stock Entry Detail')
        item_code_field =  Case().when(
                (sed.item_code.isnull() | sed.item_code == ""), sed.original_item
            ).else_(sed.item_code             
            )
        item_code_list = {row.item_code for row in self.required_items}
        data = frappe.qb.from_(se
            ).join(sed
            ).on(se.name == sed.parent
            ).select(item_code_field, Sum(sed.qty)
            ).groupby(item_code_field
            ).where(se.docstatus == 1
            ).where(se.purpose == "Material Transfer for Manufacture"
            ).where(se.work_order == self.name
            ).where(sed.item_code.isin(item_code_list) | sed.original_item.isin(item_code_list)
            ).run()
        qty_dict = frappe._dict(data)
        for d in self.required_items:            
            d.db_set("transferred_qty", flt(qty_dict.get(d.item_code)), update_modified=False)            

    def update_consumed_qty_for_required_items(self):
        """
        Update consumed qty from submitted stock entries
        against a work order for each stock item
        """
        se = frappe.qb.DocType('Stock Entry')
        sed = frappe.qb.DocType('Stock Entry Detail')
        item_code_field =  Case().when(
                (sed.item_code.isnull() | sed.item_code == ""), sed.original_item
            ).else_(sed.item_code             
            )
        item_code_list = {row.item_code for row in self.required_items}
        data = frappe.qb.from_(se
            ).join(sed
            ).on(se.name == sed.parent
            ).select(item_code_field, Sum(sed.qty)
            ).groupby(item_code_field
            ).where(se.docstatus == 1
            ).where(
                (se.purpose.isin(["Material Consumption for Manufacture","Manufacture"]))
                & (sed.s_warehouse.notnull())
                & (se.work_order == self.name)
                & (sed.item_code.isin(item_code_list) | sed.original_item.isin(item_code_list))
            ).run()
        qty_dict = frappe._dict(data)
        for item in self.required_items:
            item.db_set("consumed_qty", flt(qty_dict.get(item.item_code)), update_modified=False)

class CustomStockEntry(StockEntry):
    """标准功能代码啥也没干，跳过它 https://github.com/frappe/erpnext/issues/34729"""
    def validate_qty(self):
        pass

class CustomPickList(PickList):
    @frappe.whitelist()
    def set_item_locations(self, save=False):
        if self.fetch_item_locations:
            super(CustomPickList,self).set_item_locations(save=save)

    def before_submit(self):
        if not self.fetch_item_locations:
            frappe.throw(_('Please make sure Auto Assign Item Location checked before submit'))
        super(CustomPickList,self).before_submit()